import torch
import torch.optim as optim
import os
from clip_models import MainBackbone, ContrastiveLoss, freeze_model, VAE, vae_loss
from utils.util import getEEGDataloader, load_picture, similarity_function
import numpy as np

# 定义训练过程        
def train(model, train_loader, optimizer, criterion, scheduler, device, epoch):
    model.train()
    train_loss = 0.0
    for image, eeg_feature, eeg_label in train_loader:
        image = image.to(device)
        eeg_label = eeg_label.to(device)
        eeg_feature = eeg_feature.to(device)
        q_image, q_eeg = model(image, eeg_feature)  # 获取图像和 EEG 特征的表示
        loss = criterion(q_image, q_eeg, eeg_label)  # 使用 q_image 和 q_eeg 计算损失

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss += loss.item()

    scheduler.step()
    epoch_loss = train_loss / len(train_loader)
    print(f'Epoch {epoch+1}, Train Loss: {epoch_loss:.4f}')
    return epoch_loss


class_num = 2
batch_size = 64
gpu_id = 3
epochs = 20
is_train = False
is_load_model = True
device = torch.device(f"cuda:{gpu_id}" if torch.cuda.is_available() else "cpu")
model = MainBackbone(input_dim=6000).to(device)
if is_load_model:
    model.load_state_dict(torch.load(os.path.join('/root/data/video_decoding', "model.pth"), map_location=device))

# 训练模型
loss_val_min = np.inf
if is_train:
    # 加载数据
    train_loader, val_loader = getEEGDataloader(class_num, batch_size, 3)
    model = freeze_model(model)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = ContrastiveLoss(batch_size, device)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)
    for epoch in range(epochs):
        train(model, train_loader, optimizer, criterion, scheduler, device, epoch)
        if epoch >= 4:
            torch.save(model.state_dict(), os.path.join('/root/data/video_decoding', "model.pth"))

# 分类任务
train_loader, test_loader = getEEGDataloader(class_num, 128)
images = load_picture(1).to(device) # 加载所有图片
similarities = []
correct_val = 0
total_val = 0
with torch.no_grad():
    image_features = model.image_encoder(images)
    image_features = model.projection_head1(image_features)
    for eeg, truth in test_loader:
        eeg = eeg.to(device)
        target = truth.to(device)
        # target = torch.where(target >= 15, 1, 0).long()
        eeg_feature = model.eeg_encoder(eeg)
        eeg_feature = model.projection_head2(eeg_feature)
        results = similarity_function(eeg_feature, image_features).to(device)
        correct_val += (results == target).sum().float()
        total_val += len(target)
    epoch_acc = correct_val / total_val
    print(f'val acc: {epoch_acc:.4f}')
